# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import gzip
import os
import pickle
import typing

from recordclass import dataobject

from coin.base.datetime_util import to_datetime
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.exchange.base.tick import FixedTickPrice

import experimental.prophet as prophet
from coin.experimental.jaewon.research.adaptive_edge import features, lt_model
from experimental.prophet.ops.linear_model import linear_regression_predict
from experimental.prophet.craps.universe import universe_all


def _model(symbol_nick,
           ref_symbol_nicks,
           linear_model,
           date,
           sampler_offsets=None,
           lt_model_param=None,
           debug=False):
  sampler_offsets = sampler_offsets or {}
  lt_model_param = lt_model_param or {}

  # Time
  ts = prophet.global_timestamp()

  # Symbols
  universe = {elem.nickname: elem for elem in universe_all(date)}
  symbol_info = universe[symbol_nick]
  target_coin = symbol_info.product.base.currency.lower()
  assert target_coin in ('btc', 'eth', 'eos')

  feed = prophet.fastfeed_coin(symbol_info.product, symbol_info.sub_req)

  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs
  ask0_p = book.estimated_ask0_price.named_ref('ask0_p')
  bid0_p = book.estimated_bid0_price.named_ref('bid0_p')

  # Features
  with prophet.namespace('feature'):
    feature_vars = features.features(target_coin, ref_symbol_nicks + [symbol_nick], date)

  # Sampler
  with prophet.namespace('sampler'):
    sampler_vars = features.samplers(target_coin, 'feature/')
    sampler_vars = {v.name: v for v in sampler_vars}

  # Trigger
  trigger = features.trigger(target_coin, symbol_nick, 'feature/')

  # Long-term model theo price
  with prophet.namespace('lt_model'):
    theo_ask0_p, theo_bid0_p = lt_model.theo_price(symbol_info, date=date, **lt_model_param)

  # Model
  with prophet.namespace('edge_model'):
    x_vars = [v for v in feature_vars if not v.name.endswith('.x')]

    buy_edge_bp = 15.
    sell_edge_bp = 15.

    for sampler_name in linear_model.keys():
      offset = sampler_offsets.get(sampler_name, 3.)
      # logging.info('Sampler %s: %f bp', sampler_name, offset)

      sampler = sampler_vars[sampler_name]
      buy_model = dict(linear_model[sampler_name]['buy']['model'].params)
      sell_model = dict(linear_model[sampler_name]['sell']['model'].params)
      with prophet.control_if(sampler):
        sampler_buy_edge_bp = linear_regression_predict(x_vars, buy_model) + float(offset)
        sampler_sell_edge_bp = linear_regression_predict(x_vars, sell_model) + float(offset)

      buy_edge_bp = prophet.cond(sampler, sampler_buy_edge_bp, buy_edge_bp)
      sell_edge_bp = prophet.cond(sampler, sampler_sell_edge_bp, sell_edge_bp)

    buy_p = (theo_bid0_p * (1. - buy_edge_bp / 10000.)).named_ref('buy_p')
    sell_p = (theo_ask0_p * (1. + sell_edge_bp / 10000.)).named_ref('sell_p')

  if debug:
    ts_10s = prophet.timer('10s')
    with prophet.control_if(ts % (10 * 10**9) == 0):
      prophet.printer('%t %s %s', ts, buy_p, sell_p)


class PassModelConfig(dataobject):
  tick: typing.Any = None
  machine: str = ''
  symbol_nick: str
  ref_symbol_nicks: typing.List[str] = []
  linear_model_filepath: str
  sampler_offsets: typing.Any = None
  pull_edge_bp: float = 2.
  bound_price: bool = True
  lt_model_param: typing.Any = None
  ref_time: typing.Any = None


class PassModel:
  def __init__(self, config, pos_getter=None, trigger_callback=None):
    self._config = config
    self._pos_getter = pos_getter
    self.trigger_callback = trigger_callback

    if self._config.ref_time is not None:
      self._ref_time = to_datetime(self._config.ref_time)
    else:
      self._ref_time = datetime.datetime.now()

    self._graph = None
    self._adapter = None
    self._init_graph()

    self._tick = None
    self._init_tick()

    self._ask0_p_var = self._graph.get_named_variable('ask0_p')
    self._bid0_p_var = self._graph.get_named_variable('bid0_p')
    self._buy_p_var = self._graph.get_named_variable('edge_model/buy_p')
    self._sell_p_var = self._graph.get_named_variable('edge_model/sell_p')
    self._timestamp = self._graph.get_named_variable('global_timestamp')
    self._trigger = self._graph.get_named_variable('trigger')

    self._buy_price = None
    self._buy_price_pull = None
    self._sell_price = None
    self._sell_price_pull = None

    self._last_triggered = 0

  def _init_tick(self):
    if self._config.tick is not None:
      tick = self._config.tick
      self._tick = FixedTickPrice(tick) if isinstance(tick, float) else tick
      return

    universe = {elem.nickname: elem for elem in universe_all(self._ref_time)}
    symbol_info = universe[self._config.symbol_nick]
    self._tick = symbol_info.product_info.gen_tick_price()
    assert isinstance(self._tick, FixedTickPrice)

  def _init_graph(self):
    linear_model_filepath = os.path.expanduser(self._config.linear_model_filepath)
    with gzip.open(linear_model_filepath, 'rb') as f:
      linear_model = pickle.load(f)

    self._graph = prophet.Graph()
    with prophet.as_default_graph(self._graph):
      _model(self._config.symbol_nick,
             self._config.ref_symbol_nicks,
             linear_model,
             self._ref_time,
             sampler_offsets=self._config.sampler_offsets,
             lt_model_param=self._config.lt_model_param)

    self._adapter = prophet.FastfeedGraphAdapter(self._graph,
                                                 self._config.machine,
                                                 eval_callback=self._on_eval)

  def _on_eval(self, graph):
    if not self.ready:
      return

    ts = self.timestamp
    if not self.triggered and ts - self._last_triggered < 10 * (10**9):
      return
    self._last_triggered = ts

    if self._config.bound_price:
      ask0_minus_one = self._tick.get_prev_price(self._ask0_p_var.value)
      bid0_plus_one = self._tick.get_next_price(self._bid0_p_var.value)
      self._buy_price = min(self._tick.floor(self._buy_p_var.value), ask0_minus_one, bid0_plus_one)
      self._sell_price = max(self._tick.ceil(self._sell_p_var.value), ask0_minus_one, bid0_plus_one)
    else:
      self._buy_price = self._tick.floor(self._buy_p_var.value)
      self._sell_price = self._tick.ceil(self._sell_p_var.value)

    self._buy_price_pull = self._buy_price * (1. + self._config.pull_edge_bp / 10000.)
    self._sell_price_pull = self._sell_price * (1. - self._config.pull_edge_bp / 10000.)

    if self.trigger_callback:
      self.trigger_callback()

  def prepare(self, feed_subsystem):
    feed_subsystem.subscribe_fastfeed(self._adapter.get_feed_sub_requests(),
                                      self._adapter.on_book_reset)

  @property
  def ready(self):
    return not self._buy_p_var.is_null and not self._sell_p_var.is_null

  @property
  def triggered(self):
    return self._trigger.value

  @property
  def timestamp(self):
    return self._timestamp.value

  @property
  def buy_price(self):
    return self._buy_price

  @property
  def sell_price(self):
    return self._sell_price

  @property
  def buy_price_pull(self):
    return self._buy_price_pull

  @property
  def sell_price_pull(self):
    return self._sell_price_pull


'''
# Uncomment to test functionality

from absl import app, flags
from tornado.ioloop import IOLoop

import coin.flow.subscriber as flow_subscriber
from coin.flow.simple_queue import SimpleQueue
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem

def main(argv):
  config = PassModelConfig(
      tick=0.01,
      symbol_nick='huobi-btcusd-quarter',
      linear_model_filepath='model_huobi-btcusd-quarter.pkl.gz',
      sampler_offsets=None,
      bound_price=True)


  model = PassModel(config)

  def on_trigger():
    print(datetime.datetime.now(), 'trigger', model._ask0_p_var.value, model._bid0_p_var.value)
  model.trigger_callback = on_trigger

  feed_sub_req = model.fastfeed_adapter.get_feed_sub_requests()

  ioloop = IOLoop.current()
  queue = SimpleQueue('queue')
  flow_sub = flow_subscriber.from_queue(queue)

  feed_subsystem = MultiFeedSubsystem()
  model.prepare(feed_subsystem)
  feed_subsystem.start(ioloop, queue, flow_sub)

  try:
    ioloop.start()
    return 0
  except KeyboardInterrupt:
    return 1


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG')
  app.run(main)
'''
